# 导入所需的库
import yfinance as yf
import pandas as pd
import numpy as np

# 获取所有A股股票列表（这里只举例三只，实际上可以用yfinance或者其他方法获取更多）
stock_list = ['600000.SS', '600004.SS', '600006.SS']


# 定义一个函数用来计算单个股票的集中度
def get_concentration(ts_code):
    # 获取股票的日线数据，复权后按日期升序排列
    df = yf.download(ts_code, start='2020-01-01', end='2020-12-31', auto_adjust=True)
    df = df.sort_index()
    # 初始化筹码分布字典，键为价格，值为筹码数量（成交量）
    chip_dict = {}
    # 遍历每一天的数据
    for i in range(len(df)):
        # 获取当天的收盘价和成交量
        price = df.iloc[i]['Close']
        vol = df.iloc[i]['Volume']
        # 如果当天没有成交量，则跳过
        if vol == 0:
            continue
        # 如果筹码分布字典为空，则直接添加当天的价格和成交量作为初始值
        if not chip_dict:
            chip_dict[price] = vol
        else:
            # 否则，按照一定比例从所有价格上卖出相同数量的筹码，并在当天价格上买入相同数量的筹码（简化算法）
            ratio = vol / sum(chip_dict.values())
            for key in list(chip_dict.keys()):
                chip_dict[key] *= (1 - ratio)
                if chip_dict[key] < 0.01:  # 如果某个价格上的筹码数量太小，则删除该键值对，避免浪费内存空间
                    del chip_dict[key]
            chip_dict[price] = chip_dict.get(price, 0) + vol

    # 计算集中度指标，如CR4、CR8、HHI等，并返回结果（这里只举例CR4和HHI）
    CR4 = 0  # 前四家企业所占市场份额之和
    HHI = 0  # 所有企业市场份额平方之和

    total_chip = sum(chip_dict.values())  # 总筹码数量

    sorted_chip_list = sorted(chip_dict.items(), key=lambda x: x[1], reverse=True)  # 按照筹码数量降序排列

    for i in range(min(4, len(sorted_chip_list))):  # 取前四个或者全部（如果不足四个）
        share = sorted_chip_list[i][1] / total_chip  # 计算市场份额
        CR4 += share
        HHI += share ** 2

    return CR4, HHI


# 定义一个空的DataFrame用来存储所有股票的集中度结果
result_df = pd.DataFrame(columns=['ts_code', 'CR4', 'HHI'])

# 遍历所有股票列表，并调用集中度函数计算结果，追加到DataFrame中
for i in range(len(stock_list)):
    ts_code = stock_list[i]
    print(f'正在计算第{i}')
    print(f'股票代码:{ts_code}')
    print(f'集中度为:{get_concentration(ts_code)}')
